package com.zh.note.leetcode;

/**
 * 698. 划分为k个相等的子集
 * 给定一个整数数组  nums 和一个正整数 k，找出是否有可能把这个数组分成 k 个非空子集，其总和都相等。
 * 输入： nums = [4, 3, 2, 3, 5, 2, 1], k = 4
 * 输出： True
 * 说明： 有可能将其分成 4 个子集（5），（1,4），（2,3），（2,3）等于总和。
 */
public class LC_698_canPartitionKSubsets {
    public static boolean canPartitionKSubsets(int[] nums, int k) {
        int sum = 0;
        for (int i = 0; i < nums.length; i++) {
            sum += nums[i];
        }
        if (sum % k != 0) {
            return false;
        }
        int[] bucket = new int[k + 1];
        int target = sum / k;
        return dfs(nums, 0, bucket, k, target);

    }

    /**
     * @param nums
     * @param index  表示选择哪一个球，从第一个球开始
     * @param bucket 存储球放在桶中的之和
     * @param k      表示子集个数（桶的数量）
     * @param target 表示每个子集目标和
     * @return
     */
    private static boolean dfs(int[] nums, int index, int[] bucket, int k, int target) {
        if (index == nums.length) { //所有元素枚举完成
            for (int i = 0; i < k; i++) { //判断每个桶是否装满
                if (bucket[i] != target) {
                    return false;
                }
            }
            return true;
        }
        for (int i = 0; i < k; i++) { //第index球放到哪个桶
            if (bucket[i] + nums[index] > target) {
                continue;
            }
            // 原因：如果元素和相等，那么 nums[index] 选择上一个桶和选择当前桶可以得到的结果是一致的
            if (i > 0 && bucket[i] == bucket[i - 1]) {
                continue;
            }


            bucket[i] += nums[index];
            if (dfs(nums, index + 1, bucket, k, target)) { // 处理下一个球，即 nums[index + 1]
                return true;
            }
            bucket[i] -= nums[index];

            // if (bucket[i] == 0) {
            //     break;
            // }
        }
        return false;


    }

    public static void main(String[] args) {
        int[] nums = {5, 2, 1, 5, 2, 1, 5, 2, 1};
        System.out.println(canPartitionKSubsets(nums, 4));

    }

}
